import torch as pt
from torch.utils import data
import numpy as np
from python_ai.common.xcommon import *
import torchvision as ptv
import sys
import os

ROOT_DIR = '../../../../large_data/DL2/pt/cifar10'
BATCH_SIZE = 8000

if not os.path.exists(ROOT_DIR):
    print(f'Root dir "{ROOT_DIR}" is not right!', file=sys.stderr)
    sys.exit(1)
# ATTENTION TypeError: default_collate: batch must contain tensors, numpy arrays, numbers, dicts or lists; found <class 'PIL.Image.Image'>
# ds = ptv.datasets.MNIST(root=ROOT_DIR, train=True, download=True)
ds = ptv.datasets.CIFAR10(root=ROOT_DIR, train=False, download=True, transform=ptv.transforms.ToTensor())
print('data', ds.data.shape)
print('targets', np.shape(ds.targets))
print('train', ds.train)

dl = data.DataLoader(ds, batch_size=BATCH_SIZE, shuffle=True, drop_last=False)

for i, (bx, by) in enumerate(dl):
    sep(i)
    bx = bx.numpy()
    by = by.numpy()
    print_numpy_ndarray_info(bx, 'bx')
    print_numpy_ndarray_info(by, 'by')
